package com.simple.llm.component;

import com.alibaba.fastjson.JSON;
import com.plexpt.chatgpt.entity.billing.Usage;
import com.plexpt.chatgpt.entity.chat.ChatChoice;
import com.plexpt.chatgpt.entity.chat.ChatCompletionResponse;
import com.plexpt.chatgpt.entity.chat.Message;
import com.plexpt.chatgpt.listener.AbstractStreamListener;
import com.simple.llm.domain.dto.LLMStreamResultDTO;
import com.simple.llm.domain.dto.StreamChatDTO;
import com.simple.llm.domain.vo.ChatVO;
import com.simple.llm.domain.vo.LLMResponseVO;
import lombok.extern.slf4j.Slf4j;
import okhttp3.sse.EventSource;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;

import java.io.IOException;
import java.util.List;
import java.util.function.Consumer;

@Slf4j
public class OpenAiInterfaceLLMStreamListener extends AbstractStreamListener {

    private final ResponseBodyEmitter responseBodyEmitter;
    private final String parentMessageId;
    private final String msgId;

    private final String chatId;
    private final Consumer<LLMStreamResultDTO> resultHandler;
    private final StreamChatDTO streamChatDTO;

    private ChatCompletionResponse lastResponse;
    //是否首次输出，用于控制换行
    private boolean first = true;

    public OpenAiInterfaceLLMStreamListener(ResponseBodyEmitter responseBodyEmitter
            , StreamChatDTO streamChatDTO
            , Consumer<LLMStreamResultDTO> resultHandler) {
        this.responseBodyEmitter = responseBodyEmitter;
        this.parentMessageId = streamChatDTO.getParentMessageId();
        this.chatId = streamChatDTO.getChatId();
        this.streamChatDTO = streamChatDTO;
        this.resultHandler = resultHandler;
        this.msgId = this.chatId;
    }

    private void sendMsgToUser(String msg, boolean isNormal) {
        ChatVO chatVO = new ChatVO()
                .setRole("assistant")
                .setParentMessageId(parentMessageId)
                .setId(isNormal ? msgId : parentMessageId)
                .setText(msg);
        String sendJson = JSON.toJSONString(chatVO);
        try {
            responseBodyEmitter.send(first ? sendJson : "\n" + sendJson);
            if (!isNormal) {
                responseBodyEmitter.complete();
            }
        } catch (IOException e) {
            log.error("发送消息异常", e);
        } finally {
            if (first) {
                first = false;
            }
        }
    }

    private void finalHandler() {
        Usage usage = lastResponse.getUsage();
        LLMResponseVO vo = LLMResponseVO.builder()
                .content(this.lastMessage)
                .inputTokenCount(usage.getPromptTokens())
                .outputTokenCount(usage.getCompletionTokens())
                .desc(this.streamChatDTO.getDesc())
                .createUser(this.streamChatDTO.getCreateUser())
                .modelName(this.streamChatDTO.getModelName())
                .build();

        LLMStreamResultDTO dto = LLMStreamResultDTO.builder()
                .chatId(this.chatId)
                .messageId(this.msgId)
                .parentMessageId(parentMessageId)
                .status(true)
                .llmResponseVO(vo)
                .build();
        this.resultHandler.accept(dto);
        this.responseBodyEmitter.complete();
    }

    @Override
    public void onEvent(EventSource eventSource, String id, String type, String data) {
        if (data.equals("[DONE]")) {
            finalHandler();
            return;
        }
        lastResponse = JSON.parseObject(data, ChatCompletionResponse.class);
        // 读取Json
        List<ChatChoice> choices = lastResponse.getChoices();
        if (choices == null || choices.isEmpty()) {
            return;
        }
        Message delta = choices.get(0).getDelta();
        String text = delta.getContent();

        if (text != null) {
            lastMessage += text;
            onMsg(text);
        }
    }

    @Override
    public void onMsg(String s) {
        sendMsgToUser(this.lastMessage, true);
    }

    @Override
    public void onError(Throwable throwable, String s) {
        log.error("gpt请求异常:" + s, throwable);
        LLMStreamResultDTO dto = LLMStreamResultDTO.builder()
                .chatId(this.chatId)
                .messageId(msgId)
                .parentMessageId(parentMessageId)
                .status(false)
                .errorMsg(s)
                .error(throwable)
                .build();
        this.resultHandler.accept(dto);
        sendMsgToUser("网络连接异常，请重试", false);
    }
}
